package data.analysis.grammar;

import data.*;
import data.analysis.Statistic;
import data.analysis.grammar.symbol.BaseScope;
import data.analysis.grammar.symbol.NativeType;
import data.analysis.grammar.symbol.Scope;
import data.analysis.grammar.symbol.ScopedSymbol;
import data.analysis.grammar.symbol.Symbol;
import data.analysis.grammar.symbol.SymbolTable;

import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.tools.Diagnostic;
import javax.tools.DiagnosticCollector;
import javax.tools.JavaFileObject;

import org.antlr.runtime.ANTLRInputStream;
import org.antlr.runtime.CharStream;
import org.antlr.runtime.RecognitionException;
import org.antlr.runtime.RuleReturnScope;
import org.antlr.runtime.TokenRewriteStream;
import org.antlr.runtime.tree.CommonTree;
import org.antlr.runtime.tree.CommonTreeNodeStream;
import org.antlr.stringtemplate.StringTemplate;
import org.antlr.stringtemplate.StringTemplateGroup;

import script.grammar.FilterTranslator;
import stream.DebugStream;
import stream.ErrorStream;

import compiler.ScriptCompiler;
import compiler.ScriptCompilerException;

/**
 * A translator for statistic scripts. Used for converting statistic script source code to 
 * compiled and instantiated {@link Statistic} objects.
 * @author Charles
 *
 */
public class StatTranslator {
	private static final LocStatASTAdaptor adaptor = new LocStatASTAdaptor();
	
	private static final String CLASS_DELIM = "#";
	private static final String PACKAGE_NAME = "data.analysis.runtime";
	private static final String CLASS_NAME_PREFIX = "StatisticImpl";
	
	private SymbolTable symbolTable;
	private OutputDictionary dictionary;
	private FilterTranslator trans;
	private Map<Scope,Class<?>> filterTypeMap;
	
	private DebugStream d;
	private ErrorStream e;
	
	/**
	 * Creates a new translator.
	 * @param d the {@link DebugLogger} to use.
	 * @param e the {@link ErrorLogger} to use.
	 */
	public StatTranslator(FilterTranslator trans, DebugStream d, ErrorStream e) {
		this.trans = trans;
		this.d = d;
		this.e = e;
		initSymbols();
	}
	
	/**
	 * Creates a list of {@link Statistic} objects from statistic script source 
	 * code.
	 * @param in a stream initialized to the beginning of the source code.
	 * @return a list of <code>Statistic</code> objects.
	 * @throws IOException if there is a problem reading from the input stream.
	 */
	public List<Statistic> buildStats(InputStream in) throws IOException {
		CommonTreeNodeStream nodes = parse(in);
	        
        if (nodes == null) {
        	return null;
        }
        
        
    	String compiled = genCode(nodes);
    	if (compiled == null) {
    		return null;
    	}
    	
    	d.write("Compiling statistics script...");
    	
    	List<Statistic> list = new ArrayList<Statistic>();
    	String[] classes = compiled.split(CLASS_DELIM);
    	final DiagnosticCollector<JavaFileObject> diags = new DiagnosticCollector<JavaFileObject>();
		ScriptCompiler<Statistic> compiler = 
			new ScriptCompiler<Statistic>(ClassLoader.getSystemClassLoader());
		for (int i = 0; i < classes.length; ++i) {
			Statistic s = genStat(i, classes[i], compiler, diags);
			if (s != null) {
				list.add(s);
			}
		}
		d.write("...done.");
    	return list;
	}
	
	protected CommonTreeNodeStream parse(InputStream in) throws IOException {
		d.write("Parsing statistics script source...");
		CharStream input = new ANTLRInputStream(in);
        LocStatLexer lex = new LocStatLexer(input);
        final TokenRewriteStream tokens = new TokenRewriteStream(lex);
        LocStatParser p = new LocStatParser(tokens);
        p.setTreeAdaptor(adaptor);
        RuleReturnScope r = null;
        
		try {
			r = p.start();
			if (lex.getNumberOfSyntaxErrors() + p.getNumberOfSyntaxErrors() > 0) {
				return null;
			}
			d.write("...AST generated...running semantic analysis...");
			
			CommonTree t = (CommonTree)r.getTree();
	        CommonTreeNodeStream nodes = new CommonTreeNodeStream(t);
	        nodes.setTokenStream(tokens);
	        nodes.setTreeAdaptor(adaptor);
	        
	        LocStatCheckPass pass = new LocStatCheckPass(nodes, symbolTable);
	        e.resetCount();
	        pass.downup(t);
	        if (e.getCount() == 0) {
	        	d.write("...done.");
	        	return nodes;
	        } else {
	        	return null;
	        }
		} catch (RecognitionException e1) {
			e1.printStackTrace();
			return null;
		}
	}
	
	
	protected String genCode(CommonTreeNodeStream nodes) {
		d.write("Generating Java target code...");
		InputStream templIn = ClassLoader.getSystemClassLoader().getResourceAsStream("data/analysis/grammar/LocStatTempl.stg");
    	StringTemplateGroup templates = new StringTemplateGroup(new InputStreamReader(templIn));
    	
    	LocStatGen gen = new LocStatGen(nodes, dictionary, trans, filterTypeMap, e);
    	gen.setTemplateLib(templates);
    	StringTemplate template;
		try {
			template = (StringTemplate)gen.start().getTemplate();
		} catch (RecognitionException e1) {
			e1.printStackTrace();
			return null;
		}
    	d.write("...done:\n\n" + template + "\n");
    	return template.toString();
	}
	
	protected Statistic genStat(int i, String code, ScriptCompiler<Statistic> compiler, DiagnosticCollector<JavaFileObject> diags) {
		String qualifiedName = PACKAGE_NAME + "." + CLASS_NAME_PREFIX + i;
		d.write("...compiling " + qualifiedName + " ...");
		try {
			Class<Statistic> obj = compiler.compile(qualifiedName, code, diags, Statistic.class);
			return obj.newInstance();
		} catch (ScriptCompilerException e) {
			for (Diagnostic<? extends JavaFileObject> diag : e.getDiagnostics().getDiagnostics()) {
				this.e.write("Error: " + diag.getMessage(null));
			}
		} catch (InstantiationException e) {
			this.e.write("Error: Could not instantiate compiled statistics script object " + i + ".");
		} catch (IllegalAccessException e) {
			this.e.write("Error: Could not instantiate compiled statistics script object " + i + ".");
		}
		return null;
	}
	
	private void initSymbols() {
		Scope global = new BaseScope("Global Scope");
		
		ScopedSymbol cell = new ScopedSymbol("cell", NativeType.INT, global);
		ScopedSymbol foci = new ScopedSymbol("foci", NativeType.INT, global);
		ScopedSymbol lbl = new ScopedSymbol("lbl", NativeType.INT, global);
		ScopedSymbol clstr = new ScopedSymbol("clstr", NativeType.INT, global);
		Symbol scr = new Symbol("script", NativeType.BOOLEAN, global);
		
		cell.define(foci);
		cell.define(clstr);
		cell.define(scr);
		
		Symbol fociCyt = new Symbol("cyt", NativeType.BOOLEAN, foci);
		Symbol fociMemb = new Symbol("memb", NativeType.BOOLEAN, foci);
		Symbol clstrCyt = new Symbol("cyt", NativeType.BOOLEAN, clstr);
		Symbol clstrMemb = new Symbol("memb", NativeType.BOOLEAN, clstr);
		
		foci.define(fociCyt);
		foci.define(fociMemb);
		foci.define(lbl);
		foci.define(scr);
		clstr.define(clstrCyt);
		clstr.define(clstrMemb);
		clstr.define(lbl);
		clstr.define(scr);
		
		global.define(cell);
		global.define(foci);
		global.define(clstr);
		
		symbolTable = new SymbolTable(global, d, e);
		
		dictionary = new OutputDictionary();
		dictionary.put(global, cell, "Cell", "((Analyzer.Results)arg).getCells()");
		dictionary.put(global, foci, "Focus", "((Analyzer.Results)arg).getFoci()");
		dictionary.put(global, clstr, "Cluster", "((Analyzer.Results)arg).getClusters()");
		
		dictionary.put(cell, foci, "Focus", "((Cell)arg).getFoci()");
		dictionary.put(cell, clstr, "Cluster", "((Cell)arg).getClusters()");
		
		dictionary.put(foci, fociCyt, "boolean", "!((Focus)arg).getCluster().isMembraneAssoc()");
		dictionary.put(foci, fociMemb, "boolean", "((Focus)arg).getCluster().isMembraneAssoc()");
		dictionary.put(foci, lbl, "Object", "java.util.Collections.nCopies(((Focus)arg).getCluster().getNumLabels(), new Object())");
		dictionary.put(clstr, clstrCyt, "boolean", "!((Cluster)arg).isMembraneAssoc()");
		dictionary.put(clstr, clstrMemb, "boolean", "((Cluster)arg).isMembraneAssoc()");
		dictionary.put(clstr, lbl, "Object", "java.util.Collections.nCopies(((Cluster)arg).getNumLabels(), new Object())");
		
		filterTypeMap = new HashMap<Scope,Class<?>>();
		filterTypeMap.put(cell, Cell.class);
		filterTypeMap.put(foci, Focus.class);
		filterTypeMap.put(clstr, Cluster.class);
	}
}
